import cv2
import asyncio
import websockets
import json
import numpy as np
from ar_processor import ARProcessor
import base64
from threading import Thread


class ARServer:
    def __init__(self, marker_folder, camera_index=0):
        """初始化AR服务器"""
        # 初始化AR处理器
        self.latest_frame = None
        self.ar_processor = ARProcessor(marker_folder)
        
        # 初始化摄像头
        self.cap = cv2.VideoCapture(camera_index)
        if not self.cap.isOpened():
            raise Exception("无法打开摄像头")
            
        # WebSocket连接列表
        self.connections = set()
        
        # 最新的变换矩阵
        self.latest_transformation = None
        
        # 运行标志
        self.running = True
        
        # 启动视频处理线程
        self.video_thread = Thread(target=self.process_video)
        self.video_thread.start()

    def process_video(self):
        """处理视频流并更新变换矩阵"""
        while self.running:
            ret, frame = self.cap.read()
            if not ret:
                break

            # 处理帧获取标记信息
            result = self.ar_processor.process_frame(frame)
            
            # 解包结果：支持两种格式：
            # 1. 返回 {name: matrix}, frame
            # 2. 返回 matrix, frame
            transformation_matrix = None
            processed_frame = frame

            if isinstance(result, tuple):
                # 情况1: 返回 (matrix, frame) 或 ({name: matrix}, frame)
                first_item = result[0]
                if isinstance(first_item, dict):
                    # 如果是 {name: matrix} 字典，则取第一个矩阵
                    transformation_matrix = next(iter(first_item.values()))
                elif isinstance(first_item, np.ndarray):
                    # 如果是直接返回变换矩阵
                    transformation_matrix = first_item
                processed_frame = result[1]
            elif isinstance(result, np.ndarray):
                # 情况2: 直接返回变换矩阵
                transformation_matrix = result

            if transformation_matrix is not None:
                self.latest_transformation = transformation_matrix
            else:
                print("无法解析变换矩阵")

            # 编码帧为JPEG，降低分辨率和质量
            _, buffer = cv2.imencode(
                '.jpg',
                cv2.resize(processed_frame, (640, 480)),  # 显式设置分辨率
                [int(cv2.IMWRITE_JPEG_QUALITY), 70]  # 降低质量
            )
            
            self.latest_frame = base64.b64encode(buffer).decode('utf-8')
            
            # 小延迟以降低CPU占用
            cv2.waitKey(1)

    async def register(self, websocket):
        """注册新的WebSocket连接"""
        self.connections.add(websocket)
        try:
            await websocket.wait_closed()
        finally:
            self.connections.remove(websocket)

    async def send_data(self):
        """向所有连接的客户端发送数据"""
        while self.running:
            tasks = []
            
            if self.connections and self.latest_frame is not None:
                # 构建视频帧数据字典
                data = {
                    'frame': self.latest_frame
                }
                
                tasks = [asyncio.create_task(conn.send(json.dumps(data))) for conn in self.connections]

            if self.connections and self.latest_transformation is not None:
                # 转换变换矩阵为列表以便JSON序列化
                matrix_list = self.latest_transformation.tolist()
                
                # 构建变换矩阵数据字典
                matrix_data = {
                    'transformation_matrix': matrix_list
                }
                
                tasks.extend([asyncio.create_task(conn.send(json.dumps(matrix_data))) for conn in self.connections])

            if tasks:
                await asyncio.wait(tasks)

            # 控制最大发送频率
            await asyncio.sleep(0.033)  # 约30fps上限

    async def start_server(self, host='0.0.0.0', port=8083):
        """启动WebSocket服务器"""
        async with websockets.serve(self.register, host, port):
            await self.send_data()

    def stop(self):
        """停止服务器和视频处理"""
        self.running = False
        self.video_thread.join()
        self.cap.release()
        cv2.destroyAllWindows()


if __name__ == '__main__':
    server = None
    try:
        # 使用image文件夹中的taishili.png作为标记图像
        server = ARServer(marker_folder='image')
        print(f"AR服务器已启动，WebSocket地址: ws://localhost:8083")
        asyncio.run(server.start_server())
    except KeyboardInterrupt:
        print("检测到用户中断，正在关闭服务器...")
    except Exception as e:
        print(f"服务器启动失败: {e}")
    finally:
        if server:
            server.stop()